'''
Created on Jul 7, 2010

@author: oabalbin
'''
import sys
import numpy as np
from optparse import OptionParser
from datetime import datetime
from collections import defaultdict, deque

import RNAseq.arraytools.armanipulation as ca
import RNAseq.arraytools.subset_genes as sg
import RNAseq.io.outfile_types as ot
import RNAseq.sampleinfo.sample_info as si
import RNAseq.array as seqarray

if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--expmatrix", dest="expmatrix",
                            help="annotation file for all files to use")
    optionparser.add_option("-d", "--outfolder", dest="outfolder",
                            help="output folder")
    optionparser.add_option("-s", "--samplesFile", dest="samplesFile",
                            help="annotation file with the sample information")
    optionparser.add_option("-g", "--genesFolder", dest="genesFolder", default="use_all_genes",
                            help="annotation file with genes of interest")
    optionparser.add_option("-b", "--samplesSubset", dest="samplesSubset", default="use_all_samples",
                            help="annotation file with genes of interest")

    optionparser.add_option("--log2", action="store_true",dest="log2", default=False,
                            help="apply log 2 to the expression values")
    optionparser.add_option("--counts", action="store_true",dest="counts", default=False,
                            help="transform the RPKM to counts expression data")
    optionparser.add_option("--norm-normals", action="store_true",dest="norm_normals", default=False,
                            help="normalize by the expression of the normal samples")
    optionparser.add_option("--scale", action="store_true",dest="scale", default=False,
                            help="normalize the expression values to mean 0 and std 1")
    optionparser.add_option("--not-isoforms", action="store_true",dest="not_isoforms", default=False,
                            help="not report gene isoforms")
    optionparser.add_option("--gct-format", action="store_true",dest="gct_format", default=False,
                            help="report expression matrix in the gct format required for most gene pattern modules")
    optionparser.add_option("-t", "--tissueName", dest="tissueName",
                            help="name of the tissue analyzed")

   
    (options, args) = optionparser.parse_args()
    
    
    t = datetime.now()
    tstamp = t.strftime("%Y_%m_%d_%H_%M")
    
    outfiles = ['outfile_expression_matrix']
    time_stamp_outfiles=[]
    for name in outfiles: 
        time_stamp_outfiles.append( options.outfolder+tstamp+'_'+options.tissueName+'_'+name+'.txt')
        
    thisfile = open(options.expmatrix)
    outfile_matrix= open(time_stamp_outfiles[0],'w')
    
    ####################
    ##### Parameters    
    #  Processing parameters
    #expression_cutoff=0.0       # minimum median RPKM of expression
    median_shift=1.0
    expr_and_std_cutoff=0        #25
    not_report_isoforms=options.not_isoforms    # Do not report gene isoforms
    gct_format = options.gct_format
    
    
    ########### Read Data
    # create the meme array
    filename=options.expmatrix
    matrixfilename=options.expmatrix+'memmap'
    header_cols=2
    matobj = seqarray.read_matrix(filename, matrixfilename, header_cols)
    all_samples, all_genes, all_gene_intervals = \
    seqarray.create_matobj_dicts(matobj.cols,matobj.sample_names)
    print all_samples
    
    # sample and gene list
    '''
    select_samples_by = ['Cell Line',options.tissueName,'PASS'] #options.sampleType.replace('_',' ')
    sampleInfoFile = open(options.samplesFile)
    #sample_info_dict = si.get_sample_info(sampleInfoFile,select_samples_by)
    sample_info_dict = si.get_sample_info(sampleInfoFile)
    sample_dict = si.generate_sample_lists(sample_info_dict)
    sample_subset, sample_cols, normal_sample_cols = si.get_sample_columns(sample_dict, all_samples)
    '''
    #Subset of Samples
    if options.samplesSubset!='use_all_samples':
        
        select_samples_by = ['Cell Line',options.tissueName,'PASS'] #options.sampleType.replace('_',' ')
        sampleInfoFile = open(options.samplesFile)
        #sample_info_dict = si.get_sample_info(sampleInfoFile,select_samples_by)
        sample_info_dict = si.get_sample_info(sampleInfoFile)
        sample_dict = si.generate_sample_lists(sample_info_dict)
        sample_subset, sample_cols, normal_sample_cols = si.get_sample_columns(sample_dict, all_samples)
        
        thsamples, thnormals = sg.read_sample_list(open(options.samplesSubset))
        sample_subset, sample_cols, normal_sample_cols = sg.get_sample_subset(thsamples, thnormals, all_samples)
        print sample_subset, sample_cols, normal_sample_cols
    
    else:
        select_samples_by = ['Cell Line',options.tissueName,'PASS'] #options.sampleType.replace('_',' ')
        sampleInfoFile = open(options.samplesFile)
        sample_info_dict = si.get_sample_info(sampleInfoFile,select_samples_by)
        sample_dict = si.generate_sample_lists(sample_info_dict)
        sample_subset, sample_cols, normal_sample_cols = si.get_sample_columns(sample_dict, all_samples)

            
    if options.genesFolder!='use_all_genes':
        gene_list_folder = options.genesFolder
        # Get all the gene files in the folder
        thosefiles = sg.read_files_folder1(gene_list_folder,'.txt')
        use_all_genes=False
    else:
        use_all_genes=True
        
    ########### Program
    
    if not use_all_genes:        
        all_gene_list_names=deque()
        for gene_list_file in thosefiles:
            gene_list_name = gene_list_file.split('/')[-1]
            all_gene_list_names.append(gene_list_name)
            
            # subset of genes and samples to consider
            thisgenelist = ca.list_of_names(open(gene_list_file))
            print thisgenelist
            #print thisgenelist
            gene_subset, gene_rows, gene_subset_intervals, gene_subset_rows = \
             sg.get_gene_subset(thisgenelist, all_genes, all_gene_intervals) 
                             
            # create a sub array  using the genelist and the samplelist            
            expression_matrix = matobj.exprs[gene_rows, :]
            expression_matrix = expression_matrix[:,sample_cols]
                        
            thisGeneArray = sg.create_thisGenearray(expression_matrix, gene_subset, sample_subset, \
                                                    list(normal_sample_cols), gene_subset_intervals)
    else:
            # Get all genes from the expression matrix with the selected columns
            thisGeneArray = sg.create_thisGenearray(matobj.exprs[:,sample_cols], all_genes, sample_subset, \
                                                    list(normal_sample_cols), all_gene_intervals)
    
    # Current used transformations
    if options.counts:
        ###### Put all this into a function
        library_file = '/data/projects/libraryquality/prostate/prostate_tissue_read_duplication_scores.txt'
        library_statistics = si.read_library_statistics(library_file)
        library_size_dict = si.sample_library_statistic(library_statistics['ESTIMATED_LIBRARY_SIZE'])
        
        library_size = np.empty(len(library_size_dict), dtype=float)
        library_size_expmask=np.empty(len(library_size_dict), dtype=bool)
        
        print library_size_dict
        sps=deque()
        for i, sp in sample_subset.iteritems():            
            try:
                sps.append(sp)
                if library_size_dict[sp]=='nan':
                    library_size[i] = np.nan
                    library_size_expmask[i] = True
                else:
                    library_size[i] = float(library_size_dict[sp])
                    library_size_expmask[i] = False
                    
            except KeyError:
                print sp
                library_size[i] = np.nan
                library_size_expmask[i] = True
        #genelength = thisGeneArray.get_feature_length()
        library_size_data = np.ma.array(library_size,mask=library_size_expmask)
        libsize_mean = np.ma.mean(library_size_data)
        library_size_data = np.ma.filled(library_size_data, fill_value=libsize_mean)        
        ################################################
        
        featlength = thisGeneArray.get_feature_length()
        countmat=thisGeneArray.convert2count(thisGeneArray.expVal, featlength, library_size_data)
        thisGeneArray.expVal=countmat
        
                
    if options.log2:
        thisGeneArray.log2_norm(median_shift)
    if options.norm_normals:
        thisGeneArray.normalize_exp_bynormals(normal_sample_cols)
    if options.scale:
        thisGeneArray.center_scale_norm()
    #thisGeneArray.fold_change_accross_samples_median()

    ############# Print expression matrix
    #ca.print_expression_matrix(thisGeneArray, outfile_matrix)
    
    if not_report_isoforms:
        #thisGeneArray.average_isoforms(gene_isoforms)
        ot.print_expression_matrix_isoforms_stdfiltered(thisGeneArray, outfile_matrix, \
                                                        expr_and_std_cutoff, gct_format) # it can include the isoform list
    else:
        ot.print_expression_matrix_stdfiltered(thisGeneArray, outfile_matrix, expr_and_std_cutoff, gct_format)

        
